"""
3.	卷积神经网络常用模型有：LeNet5、AlexNet8、VGG16、InceptionNet、ResNet等，从其中任选一个模型，进行Cifar10数据集训练和分类。
利用tensorflow2.0或pytorch深度学习平台，按照下述要求，完成代码编程和演示（28分）

【Note: for Tensorflow 2.x】
"""
import tensorflow as tf
import numpy as np
from tensorflow import keras
from tensorflow.keras import layers, activations, losses, optimizers, metrics, callbacks
import os
import matplotlib.pyplot as plt

np.random.seed(777)
tf.random.set_seed(777)

ALPHA = 0.001
BATCH_SIZE = 64
N_EPOCHS = 2
VER = 'v1.0'  # 【训练后增加版本号来强制重新训练，不然训练后就直接读取保存的权重。】
FILE_NAME = os.path.basename(__file__)
LOG_DIR = os.path.join('_log', FILE_NAME, VER)
SAVE_DIR = os.path.join('_save', FILE_NAME, VER)
SAVE_PREFIX = os.path.join(SAVE_DIR, 'weights')
os.makedirs(SAVE_DIR, exist_ok=True)

# ①	导入Cifar10数据集，进行必要的预处理
(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()
M_TRAIN = len(x_train)
IMG_SHAPE = x_train[0].shape
N_CLS = len(np.unique(y_train))
rand_idx = np.random.permutation(M_TRAIN)
x_train = x_train[rand_idx]
y_train = y_train[rand_idx]
print('x_train', x_train.shape)
print('y_train', y_train.shape)
print('x_test', x_test.shape)
print('y_test', y_test.shape)

# ②	定义卷积单元convCell，由卷积、BN层、Relu激活组成
def convCell(filters, kernel_size, strides, padding='valid'):
    return keras.Sequential([
        layers.Conv2D(filters, kernel_size, strides, padding),
        layers.BatchNormalization(),
        layers.ReLU()
    ])

# ③	选取CNN模型，利用训练集数据进行模型训练
inputs = keras.Input(IMG_SHAPE)
x = convCell(16, (3, 3), (1, 1), 'same')(inputs)
x = layers.MaxPooling2D((2, 2), (2, 2), 'same')(x)
x = convCell(32, (3, 3), (1, 1), 'same')(x)
x = layers.MaxPooling2D((2, 2), (2, 2), 'same')(x)
x = layers.Flatten()(x)
x = layers.Dense(512, activation=activations.relu)(x)
x = layers.Dense(128, activation=activations.relu)(x)
x = layers.Dense(N_CLS, activation=activations.softmax)(x)
model = keras.Model(inputs, x)
model.summary()
model.compile(
    loss=losses.SparseCategoricalCrossentropy(from_logits=False),
    optimizer=optimizers.Adam(learning_rate=ALPHA),
    metrics=metrics.sparse_categorical_accuracy,
)

if len(os.listdir(SAVE_DIR)) > 0:
    print('Loading weights ...')
    model.load_weights(SAVE_PREFIX)
    print('Loaded.')
else:
    print('Training ...')
    model.fit(x_train, y_train,
              BATCH_SIZE, N_EPOCHS,
              validation_split=0.1,
              validation_batch_size=BATCH_SIZE,
              callbacks=callbacks.TensorBoard(log_dir=LOG_DIR, update_freq='batch', profile_batch=0)
              )
    print('Trained, saving ...')
    model.save_weights(SAVE_PREFIX)
    print('Saved')

# ④	模型训练完成后，利用测试集进行模型评估，输出准确率
print('Testing ...')
model.evaluate(x_test, y_test, BATCH_SIZE)
print('Tested')

# ⑤	在测试集随机选取20张图片，图形显示预测种类和标签种类（4行5列）
M_TEST = len(x_test)
rand_idx = np.random.permutation(M_TEST)
x_test = x_test[rand_idx]
y_test = y_test[rand_idx]

spr = 4
spc = 5
spn = 0
plt.figure(figsize=[8, 6])

x_test = x_test[:spr * spc]
y_test = y_test[:spr * spc]

pred = model(x_test).numpy()
pred = pred.argmax(axis=1)

for i in range(spr * spc):
    spn += 1
    plt.subplot(spr, spc, spn)
    pred_i = pred[i]
    y_i = y_test[i, 0]
    sign = 'V' if pred_i == y_i else 'X'
    plt.title(f'{y_i}: {pred_i} ({sign})')
    plt.axis('off')
    plt.imshow(x_test[i])
plt.show()
